# coding: utf-8

from PIL import Image
Image.MAX_IMAGE_PIXELS = None
import torch.nn.functional as F
import numpy as np
import shutil
import os, glob


def pad_image(image):
    interval = 2
    H, W, _ = image.shape
    pad_w, pad_h = 0, 0
    if H % interval != 0:
        pad_h = interval * (H // interval + 1) - H
    if W % interval != 0:
        pad_w = interval * (W // interval + 1) - W
    pad_image = np.pad(image, ((0, pad_h), (0, pad_w), (0, 0)), mode='constant', constant_values=0)
    H, W, _ = pad_image.shape
    return pad_image, H, W

class HM4Image:
    """
    step 1. pad image (h%2==0 and w%2==0), then convert to input.yuv.
    step 2. input.yuv to output.bin
    step 3. output.bin to output.yuv, then convert to output.png
    """
    def __init__(self,
                 input_image_folder,
                 bits_output_folder,
                 decode_output_folder,
                 qp,
                 encoder_intra_main_cfg: str = r"F:\PCL\month\2025\25-06\experiments\HM16.26\image\cfg\encoder_intra_main.cfg",
                 image_cfg: str = r"F:\PCL\month\2025\25-06\experiments\HM16.26\image\cfg\image.cfg",
                 suffix: str = ".png",
                 ffmpeg: str = r"D:\newApp\ffmpeg-7.0.2-essentials_build\bin\ffmpeg.exe",
                 encoder_cmd: str = r"D:\newApp\HM-16.8\bin\vc2015\x64\Debug\TAppEncoder.exe",
                 decoder_cmd: str = r"D:\newApp\HM-16.8\bin\vc2015\x64\Debug\TAppDecoder.exe"
                 ):
        self.suffix = suffix
        self.encoder_intra_main_cfg = encoder_intra_main_cfg
        self.image_cfg = image_cfg
        self.qp = qp
        self.input_image_folder = input_image_folder
        self.bits_output_folder = os.path.join(bits_output_folder, str(qp))    # 存放 bin 文件
        self.decode_output_folder = os.path.join(decode_output_folder, str(qp))    # 存放解码后的 png
        os.makedirs(self.bits_output_folder, exist_ok=True)
        os.makedirs(self.decode_output_folder, exist_ok=True)
        self.ffmpeg = ffmpeg
        self.encoder_cmd = encoder_cmd
        self.decoder_cmd = decoder_cmd
        self.image_names = [i for i in os.listdir(self.input_image_folder) if suffix in i]

    def start(self):
        for image_name in self.image_names:
            # step 1. pad image to mod 2 equal 0
            input_image = np.array(Image.open(os.path.join(self.input_image_folder, image_name)))
            origin_height, origin_width, _ = input_image.shape
            input_pad_image, pad_image_h, pad_image_w = pad_image(input_image)
            tmp_save_path = os.path.join(self.bits_output_folder, image_name)
            Image.fromarray(input_pad_image).save(tmp_save_path)
            # step 2. convet *.png to *.yuv format
            yuv_save_path = tmp_save_path.replace(self.suffix, ".yuv")
            cmd1 = f"{self.ffmpeg} -i {tmp_save_path} -s {pad_image_w}x{pad_image_h} -pix_fmt yuv420p {yuv_save_path}"
            os.system(cmd1)
            shutil.os.remove(tmp_save_path)
            # step 3. *.yuv to *.bin
            bin_save_path = os.path.join(self.bits_output_folder, image_name.replace(self.suffix, ".bin"))
            cmd2 = f"{self.encoder_cmd} -c {self.encoder_intra_main_cfg} -c {self.image_cfg} -b {bin_save_path} -q {self.qp} -wdt {pad_image_w} -hgt {pad_image_h} -i {yuv_save_path} -o {bin_save_path.replace('.bin', '_enc.yuv')}"
            os.system(cmd2)
            shutil.os.remove(yuv_save_path)
            # step 4. *.bin to *.yuv
            yuv_output_path = os.path.join(self.decode_output_folder, image_name.replace(self.suffix, ".yuv"))
            cmd3 = f"{self.decoder_cmd} -b {bin_save_path} -o {yuv_output_path}"
            os.system(cmd3)
            # step 5. *.yuv to *.png
            output_png_path = yuv_output_path.replace(".yuv", "_dec"+self.suffix)
            cmd4 = f"{self.ffmpeg} -f rawvideo -pix_fmt yuv420p -s {pad_image_w}x{pad_image_h} -i {yuv_output_path} -frames:v 1 {output_png_path}"
            os.system(cmd4)
            shutil.os.remove(yuv_output_path)
            # step 6. correct *png shape
            rec_image = np.array(Image.open(output_png_path))
            Image.fromarray(rec_image[:origin_height, :origin_width, :]).save(output_png_path)
            print(f"Processed {image_name}.")

    def single_image(self, image_name):
        # step 1. pad image to mod 2 equal 0
        input_image = np.array(Image.open(os.path.join(self.input_image_folder, image_name)))
        origin_height, origin_width, _ = input_image.shape
        input_pad_image, pad_image_h, pad_image_w = pad_image(input_image)
        tmp_save_path = os.path.join(self.bits_output_folder, image_name)
        Image.fromarray(input_pad_image).save(tmp_save_path)
        # step 2. convet *.png to *.yuv format
        yuv_save_path = tmp_save_path.replace(self.suffix, ".yuv")
        cmd1 = f"{self.ffmpeg} -i {tmp_save_path} -s {pad_image_w}x{pad_image_h} -pix_fmt yuv420p {yuv_save_path}"
        os.system(cmd1)
        shutil.os.remove(tmp_save_path)
        # step 3. *.yuv to *.bin
        bin_save_path = os.path.join(self.bits_output_folder, image_name.replace(self.suffix, ".bin"))
        cmd2 = f"{self.encoder_cmd} -c {self.encoder_intra_main_cfg} -c {self.image_cfg} -b {bin_save_path} -q {self.qp} -wdt {pad_image_w} -hgt {pad_image_h} -i {yuv_save_path} -o {bin_save_path.replace('.bin', '_enc.yuv')}"
        os.system(cmd2)
        shutil.os.remove(yuv_save_path)
        # step 4. *.bin to *.yuv
        yuv_output_path = os.path.join(self.decode_output_folder, image_name.replace(self.suffix, ".yuv"))
        cmd3 = f"{self.decoder_cmd} -b {bin_save_path} -o {yuv_output_path}"
        os.system(cmd3)
        # step 5. *.yuv to *.png
        output_png_path = yuv_output_path.replace(".yuv", "_dec" + self.suffix)
        cmd4 = f"{self.ffmpeg} -f rawvideo -pix_fmt yuv420p -s {pad_image_w}x{pad_image_h} -i {yuv_output_path} -frames:v 1 {output_png_path}"
        os.system(cmd4)
        shutil.os.remove(yuv_output_path)
        # step 6. correct *png shape
        rec_image = np.array(Image.open(output_png_path))
        Image.fromarray(rec_image[:origin_height, :origin_width, :]).save(output_png_path)
        print(f"Processed {image_name}.")


if __name__ == '__main__':
    qps = list(range(15, 52, 2))
    input_image_folders = [
        r"D:\video-communication-dataset\图像编解码测试\CLIC\professional_valid",
        r"D:\video-communication-dataset\图像编解码测试\kodak\archive",
        r"D:\video-communication-dataset\图像编解码测试\Tecnick\TESTIMAGES\RGB\RGB_OR_1200x1200"
    ]
    bits_output_folders = [
        r"F:\PCL\month\2025\25-06\experiments\HM16.26\image\CLIC_encode_bits",
        r"F:\PCL\month\2025\25-06\experiments\HM16.26\image\Kodak_encode_bits",
        r"F:\PCL\month\2025\25-06\experiments\HM16.26\image\Tecnick_encode_bits"
    ]
    decode_output_folders = [
        r"F:\PCL\month\2025\25-06\experiments\HM16.26\image\CLIC_decode",
        r"F:\PCL\month\2025\25-06\experiments\HM16.26\image\Kodak_decode",
        r"F:\PCL\month\2025\25-06\experiments\HM16.26\image\Tecnick_decode"
    ]
    for i in range(3):
        for qp in qps:
            hm_codec = HM4Image(
                input_image_folder=input_image_folders[i],
                bits_output_folder=bits_output_folders[i],
                decode_output_folder=decode_output_folders[i],
                qp=qp
            )
            hm_codec.start()
